package main

import "fmt"

// TailCallFrame represents a tail call that needs to be executed
type TailCallFrame struct {
	Function *Function
	Args     []*Value
}

// EvalResult represents the result of evaluation, which can be either a value or a tail call
type EvalResult struct {
	Value    *Value
	TailCall *TailCallFrame
	Error    error
}

// NewValueResult creates a result containing a value
func NewValueResult(value *Value) *EvalResult {
	if value == nil {
		value = NewNil() // Use a proper nil value instead of Go nil
	}
	return &EvalResult{Value: value}
}

// NewTailCallResult creates a result containing a tail call
func NewTailCallResult(fn *Function, args []*Value) *EvalResult {
	return &EvalResult{TailCall: &TailCallFrame{Function: fn, Args: args}}
}

// NewErrorResult creates a result containing an error
func NewErrorResult(err error) *EvalResult {
	return &EvalResult{Error: err}
}

// EvalWithTCO evaluates a Scheme expression with tail call optimization
func EvalWithTCO(expr *Value, env *Environment) (*Value, error) {
	// 添加defer/recover机制来捕获和处理continuation panic
	var result *Value
	var err error

	// 使用匿名函数和defer/recover来处理continuation panic
	evalFn := func() {
		defer func() {
			if r := recover(); r != nil {
				// 检查是否是continuation panic
				if cont, ok := r.(Continuation); ok {
					// 如果是continuation panic，设置结果为continuation的值
					result = cont.Value
					err = nil
				} else {
					// 其他panic，重新抛出
					panic(r)
				}
			}
		}()

		// 执行正常的求值过程
		internalResult := evalInternal(expr, env, true)

		// Handle tail call optimization loop
		for internalResult.TailCall != nil {
			if internalResult.Error != nil {
				err = internalResult.Error
				return
			}

			tailCall := internalResult.TailCall
			fn := tailCall.Function
			args := tailCall.Args

			// Create new environment for function scope
			fnEnv := NewEnvironment(fn.Closure)

			// Bind parameters to arguments
			for i, param := range fn.Parameters {
				fnEnv.Define(param, args[i])
			}

			// Evaluate function body with tail call optimization
			internalResult = evalInternal(fn.Body, fnEnv, true)
		}

		// 如果没有tail call了，设置结果或错误
		if internalResult.Error != nil {
			err = internalResult.Error
		} else {
			result = internalResult.Value
		}
	}

	// 调用包装后的求值函数
	evalFn()

	return result, err
}

// 注意：在Scheme中，continuation的正确实现需要保存和恢复完整的执行上下文
// 当前的实现使用panic/recover机制来模拟非本地控制流
// 当continuation被调用时，它会直接返回参数值作为call/cc表达式的值
// 这是一个简化的实现，完整的实现需要更复杂的上下文保存和恢复机制

// evalInternal is the internal evaluation function that supports tail call optimization
func evalInternal(expr *Value, env *Environment, inTailPosition bool) *EvalResult {
	switch expr.Type {
	case TypeInteger, TypeRational, TypeReal, TypeComplex, TypeBoolean, TypeString, TypeChar, TypePair, TypeVector:
		// Self-evaluating expressions
		return NewValueResult(expr)

	case TypeSymbol:
		// Variable lookup
		value, err := env.Get(expr.Symbol)
		if err != nil {
			return NewErrorResult(err)
		}
		return NewValueResult(value)

	case TypeList:
		if len(expr.List) == 0 {
			return NewValueResult(NewNil())
		}

		// Check for special forms
		first := expr.List[0]
		if first.Type == TypeSymbol {
			switch first.Symbol {
			case "quote":
				return evalQuoteInternal(expr.List[1:], env)
			case "define":
				return evalDefineInternal(expr.List[1:], env)
			case "set!":
				return evalSetInternal(expr.List[1:], env)
			case "lambda":
				return evalLambdaInternal(expr.List[1:], env)
			case "if":
				return evalIfInternal(expr.List[1:], env, inTailPosition)
			case "cond":
				return evalCondInternal(expr.List[1:], env, inTailPosition)
			case "let":
				return evalLetInternal(expr.List[1:], env, inTailPosition)
			case "begin":
				return evalBeginInternal(expr.List[1:], env, inTailPosition)
			}
		}

		// Function application
		return evalApplicationInternal(expr.List, env, inTailPosition)

	default:
		return NewErrorResult(fmt.Errorf("cannot evaluate expression of type %v", expr.Type))
	}
}

// Tail call optimized versions of evaluation functions

// evalSetInternal handles set! special form with tail call optimization
func evalSetInternal(args []*Value, env *Environment) *EvalResult {
	if len(args) != 2 {
		return NewErrorResult(fmt.Errorf("set! expects exactly 2 arguments"))
	}

	name := args[0]
	value := args[1]

	if name.Type != TypeSymbol {
		return NewErrorResult(fmt.Errorf("set! first argument must be a symbol"))
	}

	// Evaluate the value expression
	evalValueResult := evalInternal(value, env, false)
	if evalValueResult.Error != nil {
		return evalValueResult
	}

	// Handle tail calls in value evaluation
	for evalValueResult.TailCall != nil {
		tailCall := evalValueResult.TailCall
		fnEnv := NewEnvironment(tailCall.Function.Closure)
		for i, param := range tailCall.Function.Parameters {
			fnEnv.Define(param, tailCall.Args[i])
		}
		evalValueResult = evalInternal(tailCall.Function.Body, fnEnv, true)
		if evalValueResult.Error != nil {
			return evalValueResult
		}
	}

	// Update the binding
	err := env.Set(name.Symbol, evalValueResult.Value)
	if err != nil {
		return NewErrorResult(err)
	}

	// Return the value that was set (Scheme convention)
	return NewValueResult(evalValueResult.Value)
}

func evalQuoteInternal(args []*Value, env *Environment) *EvalResult {
	if len(args) != 1 {
		return NewErrorResult(fmt.Errorf("quote expects exactly 1 argument"))
	}
	return NewValueResult(args[0])
}

func evalDefineInternal(args []*Value, env *Environment) *EvalResult {
	if len(args) != 2 {
		return NewErrorResult(fmt.Errorf("define expects exactly 2 arguments"))
	}

	first := args[0]
	second := args[1]

	// Check if first argument is a list (function definition shorthand)
	if first.Type == TypeList && len(first.List) > 0 {
		// Function definition: (define (name param1 param2 ...) body)
		functionName := first.List[0]
		if functionName.Type != TypeSymbol {
			return NewErrorResult(fmt.Errorf("function name must be a symbol"))
		}

		// Extract parameters
		var paramNames []string
		for _, param := range first.List[1:] {
			if param.Type != TypeSymbol {
				return NewErrorResult(fmt.Errorf("function parameters must be symbols"))
			}
			paramNames = append(paramNames, param.Symbol)
		}

		// Create lambda function
		lambdaFn := NewFunction(paramNames, second, env)
		env.Define(functionName.Symbol, lambdaFn)
		return NewValueResult(NewSymbol(functionName.Symbol))
	} else if first.Type == TypeSymbol {
		// Variable definition: (define name value)
		valueResult := evalInternal(second, env, false)
		if valueResult.Error != nil {
			return valueResult
		}

		env.Define(first.Symbol, valueResult.Value)
		return NewValueResult(NewSymbol(first.Symbol))
	} else {
		return NewErrorResult(fmt.Errorf("define expects a symbol or function definition as first argument"))
	}
}

func evalLambdaInternal(args []*Value, env *Environment) *EvalResult {
	if len(args) < 2 {
		return NewErrorResult(fmt.Errorf("lambda expects at least 2 arguments"))
	}

	params := args[0]
	var body *Value
	if len(args) == 2 {
		// Single expression body
		body = args[1]
	} else {
		// Multiple expressions body, wrap in begin
		body = &Value{
			Type: TypeList,
			List: append([]*Value{NewSymbol("begin")}, args[1:]...),
		}
	}

	if params.Type != TypeList {
		return NewErrorResult(fmt.Errorf("lambda expects a list of parameters"))
	}

	var paramNames []string
	for _, param := range params.List {
		if param.Type != TypeSymbol {
			return NewErrorResult(fmt.Errorf("lambda parameters must be symbols"))
		}
		paramNames = append(paramNames, param.Symbol)
	}

	return NewValueResult(NewFunction(paramNames, body, env))
}

func evalIfInternal(args []*Value, env *Environment, inTailPosition bool) *EvalResult {
	if len(args) < 2 || len(args) > 3 {
		return NewErrorResult(fmt.Errorf("if expects 2 or 3 arguments"))
	}

	conditionResult := evalInternal(args[0], env, false)
	if conditionResult.Error != nil {
		return conditionResult
	}

	// Handle tail call in condition evaluation
	for conditionResult.TailCall != nil {
		tailCall := conditionResult.TailCall
		fnEnv := NewEnvironment(tailCall.Function.Closure)
		for i, param := range tailCall.Function.Parameters {
			fnEnv.Define(param, tailCall.Args[i])
		}
		conditionResult = evalInternal(tailCall.Function.Body, fnEnv, true)
		if conditionResult.Error != nil {
			return conditionResult
		}
	}

	// Safety check for nil value
	if conditionResult.Value == nil {
		return NewErrorResult(fmt.Errorf("condition evaluation returned nil value"))
	}

	if conditionResult.Value.IsTruthy() {
		return evalInternal(args[1], env, inTailPosition)
	} else if len(args) == 3 {
		return evalInternal(args[2], env, inTailPosition)
	} else {
		return NewValueResult(NewNil())
	}
}

func evalCondInternal(args []*Value, env *Environment, inTailPosition bool) *EvalResult {
	for _, clause := range args {
		if clause.Type != TypeList || len(clause.List) != 2 {
			return NewErrorResult(fmt.Errorf("cond clauses must be lists of length 2"))
		}

		condition := clause.List[0]
		result := clause.List[1]

		// Check for else clause
		if condition.Type == TypeSymbol && condition.Symbol == "else" {
			return evalInternal(result, env, inTailPosition)
		}

		condResult := evalInternal(condition, env, false)
		if condResult.Error != nil {
			return condResult
		}

		// Handle tail call in condition evaluation
		for condResult.TailCall != nil {
			tailCall := condResult.TailCall
			fnEnv := NewEnvironment(tailCall.Function.Closure)
			for i, param := range tailCall.Function.Parameters {
				fnEnv.Define(param, tailCall.Args[i])
			}
			condResult = evalInternal(tailCall.Function.Body, fnEnv, true)
			if condResult.Error != nil {
				return condResult
			}
		}

		// Safety check for nil value
		if condResult.Value == nil {
			return NewErrorResult(fmt.Errorf("condition evaluation returned nil value"))
		}

		if condResult.Value.IsTruthy() {
			return evalInternal(result, env, inTailPosition)
		}
	}

	return NewValueResult(NewNil())
}

func evalLetInternal(args []*Value, env *Environment, inTailPosition bool) *EvalResult {
	// Let should have at least 2 arguments: bindings and body
	if len(args) < 2 {
		return NewErrorResult(fmt.Errorf("let expects at least 2 arguments"))
	}

	bindings := args[0]

	if bindings.Type != TypeList {
		return NewErrorResult(fmt.Errorf("let bindings must be a list"))
	}

	// Create new environment for let scope
	letEnv := NewEnvironment(env)

	// Process bindings
	for _, binding := range bindings.List {
		if binding.Type != TypeList || len(binding.List) != 2 {
			return NewErrorResult(fmt.Errorf("let bindings must be lists of length 2"))
		}

		name := binding.List[0]
		value := binding.List[1]

		if name.Type != TypeSymbol {
			return NewErrorResult(fmt.Errorf("let binding names must be symbols"))
		}

		evalValueResult := evalInternal(value, env, false) // Evaluate in outer environment
		if evalValueResult.Error != nil {
			return evalValueResult
		}

		letEnv.Define(name.Symbol, evalValueResult.Value)
	}

	// Process body expressions
	var result *EvalResult
	i := 1
	for i < len(args) {
		// Check if current expression is a define
		if expr := args[i]; expr.Type == TypeList && len(expr.List) > 0 && 
			expr.List[0].Type == TypeSymbol && expr.List[0].Symbol == "define" {
			// Evaluate define in let environment
			result = evalDefineInternal(expr.List[1:], letEnv)
			if result.Error != nil {
				return result
			}
			i++
		} else {
			// No more defines, start evaluating body expressions
			break
		}
	}

	// If we processed all expressions, return nil
	if i >= len(args) {
		return NewValueResult(NewNil())
	}

	// Evaluate the remaining expressions sequentially
	for ; i < len(args)-1; i++ {
		result = evalInternal(args[i], letEnv, false) // Not tail position for intermediate expressions
		if result.Error != nil {
			return result
		}
		// Handle any tail calls in intermediate results
		for result.TailCall != nil {
			tailCall := result.TailCall
			fnEnv := NewEnvironment(tailCall.Function.Closure)
			for i, param := range tailCall.Function.Parameters {
				fnEnv.Define(param, tailCall.Args[i])
			}
			result = evalInternal(tailCall.Function.Body, fnEnv, false)
			if result.Error != nil {
				return result
			}
		}
	}

	// Evaluate the last expression with the proper tail position flag
	result = evalInternal(args[i], letEnv, inTailPosition)
	return result
}

func evalBeginInternal(args []*Value, env *Environment, inTailPosition bool) *EvalResult {
	if len(args) == 0 {
		return NewValueResult(NewNil())
	}

	// Evaluate all expressions except the last one with inTailPosition = false
	for i := 0; i < len(args)-1; i++ {
		result := evalInternal(args[i], env, false)
		if result.Error != nil {
			return result
		}
	}

	// Evaluate the last expression with the original inTailPosition
	return evalInternal(args[len(args)-1], env, inTailPosition)
}

// 修改后的 evalApplicationInternal 函数
func evalApplicationInternal(exprs []*Value, env *Environment, inTailPosition bool) *EvalResult {
	if len(exprs) == 0 {
		return NewValueResult(NewNil())
	}

	// Evaluate function
	fnResult := evalInternal(exprs[0], env, false)
	if fnResult.Error != nil {
		return fnResult
	}

	// 如果函数本身是一个尾调用结果，先执行尾调用获取实际函数
	if fnResult.TailCall != nil {
		// 执行尾调用获取实际函数
		for fnResult.TailCall != nil {
			tailCall := fnResult.TailCall
			fnEnv := NewEnvironment(tailCall.Function.Closure)
			for i, param := range tailCall.Function.Parameters {
				fnEnv.Define(param, tailCall.Args[i])
			}
			fnResult = evalInternal(tailCall.Function.Body, fnEnv, true)
			if fnResult.Error != nil {
				return fnResult
			}
		}
	}

	// Evaluate arguments
	var args []*Value
	for _, arg := range exprs[1:] {
		argResult := evalInternal(arg, env, false)
		if argResult.Error != nil {
			return argResult
		}

		// 如果参数是尾调用结果，先执行尾调用获取实际值
		if argResult.TailCall != nil {
			// 执行尾调用获取实际参数值
			for argResult.TailCall != nil {
				tailCall := argResult.TailCall
				fnEnv := NewEnvironment(tailCall.Function.Closure)
				for i, param := range tailCall.Function.Parameters {
					fnEnv.Define(param, tailCall.Args[i])
				}
				argResult = evalInternal(tailCall.Function.Body, fnEnv, true)
				if argResult.Error != nil {
					return argResult
				}
			}
		}

		args = append(args, argResult.Value)
	}

	// Apply function
	switch fnResult.Value.Type {
	case TypeBuiltin:
		// 为builtin函数调用添加recover机制以处理continuation
		var result *Value
		var err error
		func() {
			defer func() {
				if r := recover(); r != nil {
					if cont, ok := r.(Continuation); ok {
						// 处理continuation panic
						result = cont.Value
					} else if e, ok := r.(error); ok {
						err = e
					} else {
						// 重新抛出其他类型的panic
						panic(r)
					}
				}
			}()
			result, err = fnResult.Value.Builtin(args, env)
		}()
		if err != nil {
			return NewErrorResult(err)
		}
		return NewValueResult(result)

	case TypeFunction:
		if len(args) != len(fnResult.Value.Function.Parameters) {
			return NewErrorResult(fmt.Errorf("function expects %d arguments, got %d",
				len(fnResult.Value.Function.Parameters), len(args)))
		}

		// If in tail position, return a tail call frame instead of recursing
		if inTailPosition {
			return NewTailCallResult(fnResult.Value.Function, args)
		}

		// Not in tail position, evaluate normally
		fnEnv := NewEnvironment(fnResult.Value.Function.Closure)

		// Bind parameters to arguments
		for i, param := range fnResult.Value.Function.Parameters {
			fnEnv.Define(param, args[i])
		}

		// Evaluate function body
		return evalInternal(fnResult.Value.Function.Body, fnEnv, true)

	default:
		return NewErrorResult(fmt.Errorf("cannot apply %v as function", fnResult.Value.Type))
	}
}
